"""
MongoDB 客户端
"""

from types import TracebackType

from pymongo import MongoClient
from pymongo.client_session import ClientSession
from pymongo.database import Database
from pymongo.write_concern import WriteConcern

from config import settings

mongo_client: MongoClient = MongoClient(settings.MONGODB.URI_WITH_REPLICA_SET)
""" MongoDB 客户端对象 """

mongo_db: Database = mongo_client[settings.MONGODB.DB_NAME]
""" MongoDB 数据库对象 """


class MongoDBTransaction:
    """ MongoDB 事务 上下文管理器 """

    def __enter__(self) -> ClientSession:
        self.session: ClientSession = mongo_client.start_session()
        self.session.start_transaction(write_concern=WriteConcern("majority"))

        return self.session

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        if exc_type is None:
            self.session.commit_transaction()

        else:
            self.session.abort_transaction()

        self.session.end_session()


if __name__ == "__main__":
    with MongoDBTransaction() as session:
        mongo_db.user.insert_one({"name": "Alen"}, session=session)
        mongo_db.user.insert_one({"name": "Ben"}, session=session)
